#include "jit_compiler.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <sys/mman.h>

// x86-64 Register encoding
typedef enum {
    RAX = 0, RCX = 1, RDX = 2, RBX = 3,
    RSP = 4, RBP = 5, RSI = 6, RDI = 7,
    R8 = 8, R9 = 9, R10 = 10, R11 = 11,
    R12 = 12, R13 = 13, R14 = 14, R15 = 15
} x86_reg_t;

// x86-64 calling convention registers
#define ARG_REG_1 RDI
#define ARG_REG_2 RSI
#define ARG_REG_3 RDX
#define ARG_REG_4 RCX
#define RETURN_REG RAX

// Code generation context
typedef struct {
    uint8_t* code_buffer;
    size_t buffer_size;
    size_t buffer_pos;
    jit_function_t* function;
} codegen_context_t;

// Forward declarations
static int emit_byte(codegen_context_t* ctx, uint8_t byte);
static int emit_word(codegen_context_t* ctx, uint16_t word);
static int emit_dword(codegen_context_t* ctx, uint32_t dword);
static int emit_qword(codegen_context_t* ctx, uint64_t qword);
static int emit_rex(codegen_context_t* ctx, bool w, bool r, bool x, bool b);
static int emit_modrm(codegen_context_t* ctx, uint8_t mod, uint8_t reg, uint8_t rm);
static int emit_sib(codegen_context_t* ctx, uint8_t scale, uint8_t index, uint8_t base);
static x86_reg_t get_register_for_value(jit_value_t* value);
static int generate_instruction(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_prologue(codegen_context_t* ctx);
static int generate_epilogue(codegen_context_t* ctx);
static int generate_add(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_sub(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_mul(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_div(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_const(codegen_context_t* ctx, jit_instruction_t* inst);
static int generate_return(codegen_context_t* ctx, jit_instruction_t* inst);

// =============================================================================
// Main Code Generation Entry Point
// =============================================================================

int jit_generate_machine_code(jit_function_t* function) {
    if (!function) return -1;
    
    // Determine target architecture and call appropriate generator
    #ifdef __x86_64__
    return jit_generate_x86_64(function);
    #elif __aarch64__
    return jit_generate_arm64(function);
    #else
    // Fallback: no code generation support
    return -1;
    #endif
}

int jit_generate_x86_64(jit_function_t* function) {
    if (!function) return -1;
    
    // Allocate initial code buffer
    size_t initial_size = 4096; // 4KB initial buffer
    uint8_t* code_buffer = malloc(initial_size);
    if (!code_buffer) return -1;
    
    codegen_context_t ctx = {
        .code_buffer = code_buffer,
        .buffer_size = initial_size,
        .buffer_pos = 0,
        .function = function
    };
    
    // Generate function prologue
    if (generate_prologue(&ctx) != 0) {
        free(code_buffer);
        return -1;
    }
    
    // Generate code for each basic block
    jit_basic_block_t* block = function->blocks;
    while (block) {
        // Store code offset for this block
        block->machine_code_offset = ctx.buffer_pos;
        
        // Generate instructions
        jit_instruction_t* inst = block->first_inst;
        while (inst) {
            if (generate_instruction(&ctx, inst) != 0) {
                free(code_buffer);
                return -1;
            }
            inst = inst->next;
        }
        
        block = block->next;
    }
    
    // Generate function epilogue
    if (generate_epilogue(&ctx) != 0) {
        free(code_buffer);
        return -1;
    }
    
    // Store generated code in function
    function->machine_code = code_buffer;
    function->machine_code_size = ctx.buffer_pos;
    
    // Make memory executable
    if (mprotect(code_buffer, ctx.buffer_pos, PROT_READ | PROT_EXEC) != 0) {
        free(code_buffer);
        return -1;
    }
    
    // Set function pointer
    function->compiled_function = (void*)code_buffer;
    
    return 0;
}

// Placeholder for ARM64 code generation
int jit_generate_arm64(jit_function_t* function) {
    if (!function) return -1;
    
    // ARM64 code generation would be implemented here
    // For now, just return error
    return -1;
}

// =============================================================================
// x86-64 Code Emission Helpers
// =============================================================================

static int emit_byte(codegen_context_t* ctx, uint8_t byte) {
    if (ctx->buffer_pos >= ctx->buffer_size) {
        // Expand buffer
        size_t new_size = ctx->buffer_size * 2;
        uint8_t* new_buffer = realloc(ctx->code_buffer, new_size);
        if (!new_buffer) return -1;
        
        ctx->code_buffer = new_buffer;
        ctx->buffer_size = new_size;
    }
    
    ctx->code_buffer[ctx->buffer_pos++] = byte;
    return 0;
}

static int emit_word(codegen_context_t* ctx, uint16_t word) {
    if (emit_byte(ctx, word & 0xFF) != 0) return -1;
    if (emit_byte(ctx, (word >> 8) & 0xFF) != 0) return -1;
    return 0;
}

static int emit_dword(codegen_context_t* ctx, uint32_t dword) {
    if (emit_word(ctx, dword & 0xFFFF) != 0) return -1;
    if (emit_word(ctx, (dword >> 16) & 0xFFFF) != 0) return -1;
    return 0;
}

static int emit_qword(codegen_context_t* ctx, uint64_t qword) {
    if (emit_dword(ctx, qword & 0xFFFFFFFF) != 0) return -1;
    if (emit_dword(ctx, (qword >> 32) & 0xFFFFFFFF) != 0) return -1;
    return 0;
}

static int emit_rex(codegen_context_t* ctx, bool w, bool r, bool x, bool b) {
    uint8_t rex = 0x40;
    if (w) rex |= 0x08; // 64-bit operand
    if (r) rex |= 0x04; // Extension of ModR/M reg field
    if (x) rex |= 0x02; // Extension of SIB index field
    if (b) rex |= 0x01; // Extension of ModR/M r/m field or SIB base field
    
    return emit_byte(ctx, rex);
}

static int emit_modrm(codegen_context_t* ctx, uint8_t mod, uint8_t reg, uint8_t rm) {
    uint8_t modrm = ((mod & 0x3) << 6) | ((reg & 0x7) << 3) | (rm & 0x7);
    return emit_byte(ctx, modrm);
}

static int emit_sib(codegen_context_t* ctx, uint8_t scale, uint8_t index, uint8_t base) {
    uint8_t sib = ((scale & 0x3) << 6) | ((index & 0x7) << 3) | (base & 0x7);
    return emit_byte(ctx, sib);
}

// =============================================================================
// Register Management
// =============================================================================

static x86_reg_t get_register_for_value(jit_value_t* value) {
    if (!value) return RAX;
    
    if (value->is_parameter) {
        // Use calling convention registers for parameters
        uint32_t param_index = value->data.uint_val;
        switch (param_index) {
            case 0: return ARG_REG_1;
            case 1: return ARG_REG_2;
            case 2: return ARG_REG_3;
            case 3: return ARG_REG_4;
            default: return RAX; // Would need stack access for more parameters
        }
    }
    
    // Use assigned register or default to RAX
    if (value->reg_id >= 0 && value->reg_id < 16) {
        return (x86_reg_t)value->reg_id;
    }
    
    return RAX;
}

// =============================================================================
// Function Prologue and Epilogue
// =============================================================================

static int generate_prologue(codegen_context_t* ctx) {
    // Standard x86-64 function prologue:
    // push rbp
    // mov rbp, rsp
    // sub rsp, stack_space (if needed)
    
    // push rbp
    if (emit_byte(ctx, 0x55) != 0) return -1;
    
    // mov rbp, rsp
    if (emit_rex(ctx, true, false, false, false) != 0) return -1; // REX.W
    if (emit_byte(ctx, 0x89) != 0) return -1; // MOV r/m64, r64
    if (emit_modrm(ctx, 0x3, RSP, RBP) != 0) return -1; // RSP -> RBP
    
    // Reserve stack space if needed
    if (ctx->function->stack_size > 0) {
        // sub rsp, immediate
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0x81) != 0) return -1; // SUB r/m64, imm32
        if (emit_modrm(ctx, 0x3, 5, RSP) != 0) return -1; // Opcode extension /5
        if (emit_dword(ctx, (uint32_t)ctx->function->stack_size) != 0) return -1;
    }
    
    return 0;
}

static int generate_epilogue(codegen_context_t* ctx) {
    // Standard x86-64 function epilogue:
    // mov rsp, rbp (or leave)
    // pop rbp
    // ret
    
    // leave (equivalent to mov rsp, rbp; pop rbp)
    if (emit_byte(ctx, 0xC9) != 0) return -1;
    
    // ret
    if (emit_byte(ctx, 0xC3) != 0) return -1;
    
    return 0;
}

// =============================================================================
// Instruction Code Generation
// =============================================================================

static int generate_instruction(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (!ctx || !inst) return -1;
    
    switch (inst->opcode) {
        case JIT_OP_ADD:
            return generate_add(ctx, inst);
            
        case JIT_OP_SUB:
            return generate_sub(ctx, inst);
            
        case JIT_OP_MUL:
            return generate_mul(ctx, inst);
            
        case JIT_OP_DIV:
            return generate_div(ctx, inst);
            
        case JIT_OP_CONST:
            return generate_const(ctx, inst);
            
        case JIT_OP_RETURN:
            return generate_return(ctx, inst);
            
        default:
            // Unsupported instruction - generate NOP for now
            return emit_byte(ctx, 0x90); // NOP
    }
}

static int generate_add(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (!inst->src1 || !inst->src2 || !inst->dest) return -1;
    
    x86_reg_t src1_reg = get_register_for_value(inst->src1);
    x86_reg_t src2_reg = get_register_for_value(inst->src2);
    x86_reg_t dest_reg = get_register_for_value(inst->dest);
    
    // For simplicity, assume dest_reg = src1_reg initially
    if (dest_reg != src1_reg) {
        // mov dest_reg, src1_reg
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0x89) != 0) return -1;
        if (emit_modrm(ctx, 0x3, src1_reg, dest_reg) != 0) return -1;
    }
    
    // add dest_reg, src2_reg
    if (emit_rex(ctx, true, false, false, false) != 0) return -1;
    if (emit_byte(ctx, 0x01) != 0) return -1; // ADD r/m64, r64
    if (emit_modrm(ctx, 0x3, src2_reg, dest_reg) != 0) return -1;
    
    return 0;
}

static int generate_sub(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (!inst->src1 || !inst->src2 || !inst->dest) return -1;
    
    x86_reg_t src1_reg = get_register_for_value(inst->src1);
    x86_reg_t src2_reg = get_register_for_value(inst->src2);
    x86_reg_t dest_reg = get_register_for_value(inst->dest);
    
    // Move src1 to dest if different
    if (dest_reg != src1_reg) {
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0x89) != 0) return -1;
        if (emit_modrm(ctx, 0x3, src1_reg, dest_reg) != 0) return -1;
    }
    
    // sub dest_reg, src2_reg
    if (emit_rex(ctx, true, false, false, false) != 0) return -1;
    if (emit_byte(ctx, 0x29) != 0) return -1; // SUB r/m64, r64
    if (emit_modrm(ctx, 0x3, src2_reg, dest_reg) != 0) return -1;
    
    return 0;
}

static int generate_mul(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (!inst->src1 || !inst->src2 || !inst->dest) return -1;
    
    // x86-64 multiplication is more complex - simplified version
    x86_reg_t src1_reg = get_register_for_value(inst->src1);
    x86_reg_t src2_reg = get_register_for_value(inst->src2);
    x86_reg_t dest_reg = get_register_for_value(inst->dest);
    
    // Move src1 to RAX (required for IMUL)
    if (src1_reg != RAX) {
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0x89) != 0) return -1;
        if (emit_modrm(ctx, 0x3, src1_reg, RAX) != 0) return -1;
    }
    
    // imul rax, src2_reg (result in RAX)
    if (emit_rex(ctx, true, false, false, false) != 0) return -1;
    if (emit_byte(ctx, 0x0F) != 0) return -1;
    if (emit_byte(ctx, 0xAF) != 0) return -1; // IMUL r64, r/m64
    if (emit_modrm(ctx, 0x3, RAX, src2_reg) != 0) return -1;
    
    // Move result to destination if different
    if (dest_reg != RAX) {
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0x89) != 0) return -1;
        if (emit_modrm(ctx, 0x3, RAX, dest_reg) != 0) return -1;
    }
    
    return 0;
}

static int generate_div(codegen_context_t* ctx, jit_instruction_t* inst) {
    // Division is complex on x86-64 - simplified placeholder
    // Real implementation would handle signed/unsigned division,
    // remainder, and register constraints properly
    return emit_byte(ctx, 0x90); // NOP for now
}

static int generate_const(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (!inst->dest) return -1;
    
    x86_reg_t dest_reg = get_register_for_value(inst->dest);
    
    if (inst->dest->is_constant) {
        // mov dest_reg, immediate
        if (emit_rex(ctx, true, false, false, false) != 0) return -1;
        if (emit_byte(ctx, 0xB8 + (dest_reg & 0x7)) != 0) return -1; // MOV r64, imm64
        
        // Emit immediate value based on type
        if (jit_type_is_integer(inst->dest->type)) {
            if (emit_qword(ctx, (uint64_t)inst->dest->data.int_val) != 0) return -1;
        } else if (inst->dest->type == JIT_TYPE_DOUBLE) {
            // For floating point, we'd need to emit the bit pattern
            union { double d; uint64_t i; } u;
            u.d = inst->dest->data.float_val;
            if (emit_qword(ctx, u.i) != 0) return -1;
        } else {
            if (emit_qword(ctx, 0) != 0) return -1; // Default to 0
        }
    }
    
    return 0;
}

static int generate_return(codegen_context_t* ctx, jit_instruction_t* inst) {
    if (inst->src1) {
        // Move return value to RAX
        x86_reg_t src_reg = get_register_for_value(inst->src1);
        if (src_reg != RETURN_REG) {
            if (emit_rex(ctx, true, false, false, false) != 0) return -1;
            if (emit_byte(ctx, 0x89) != 0) return -1;
            if (emit_modrm(ctx, 0x3, src_reg, RETURN_REG) != 0) return -1;
        }
    }
    
    // The actual return will be handled by the epilogue
    return 0;
}

// =============================================================================
// Debug and Disassembly
// =============================================================================

void jit_function_print_assembly(const jit_function_t* function) {
    if (!function || !function->machine_code) {
        printf("No machine code available\n");
        return;
    }
    
    printf("Machine code for function %s:\n", function->name);
    printf("Size: %zu bytes\n", function->machine_code_size);
    
    // Simple hex dump - real implementation would use a disassembler
    for (size_t i = 0; i < function->machine_code_size; i++) {
        if (i % 16 == 0) {
            printf("\n%04zx: ", i);
        }
        printf("%02x ", function->machine_code[i]);
    }
    printf("\n");
}

int jit_dump_machine_code(const jit_function_t* function, const char* filename) {
    if (!function || !function->machine_code || !filename) return -1;
    
    FILE* f = fopen(filename, "wb");
    if (!f) return -1;
    
    size_t written = fwrite(function->machine_code, 1, function->machine_code_size, f);
    fclose(f);
    
    return (written == function->machine_code_size) ? 0 : -1;
}